"""
频域特征提取

实现了频域特征提取功能，包括：
- 功率谱密度
- 谱质心
- 谱滚降
- 谱带宽
- MFCC特征
"""

import numpy as np
from typing import Dict, List, Optional
from loguru import logger
from scipy.fft import fft
from scipy.signal import windows
import librosa


class FrequencyDomainFeatures:
    """频域特征提取器"""
    
    def __init__(self, 
                 features: List[str] = None,
                 sample_rate: int = 100,
                 n_mfcc: int = 13,
                 n_fft: int = 512,
                 hop_length: int = 128,
                 n_mels: int = 40):
        """
        初始化频域特征提取器
        
        Args:
            features: 要提取的特征列表
            sample_rate: 采样率
            n_mfcc: MFCC系数个数
            n_fft: FFT大小
            hop_length: 帧移
            n_mels: Mel滤波器组个数
        """
        if features is None:
            features = ["power_spectral_density", "spectral_centroid", 
                       "spectral_rolloff", "spectral_bandwidth", "mfcc"]
        
        self.features = features
        self.sample_rate = sample_rate
        self.n_mfcc = n_mfcc
        self.n_fft = n_fft
        self.hop_length = hop_length
        self.n_mels = n_mels
        
        logger.info(f"初始化频域特征提取器，特征: {features}")
    
    def extract_features(self, signal_data: np.ndarray) -> Dict[str, np.ndarray]:
        """
        提取频域特征
        
        Args:
            signal_data: 输入信号
            
        Returns:
            特征字典
        """
        feature_dict = {}
        
        for feature in self.features:
            if feature == "power_spectral_density":
                feature_dict[feature] = self._extract_power_spectral_density(signal_data)
            elif feature == "spectral_centroid":
                feature_dict[feature] = self._extract_spectral_centroid(signal_data)
            elif feature == "spectral_rolloff":
                feature_dict[feature] = self._extract_spectral_rolloff(signal_data)
            elif feature == "spectral_bandwidth":
                feature_dict[feature] = self._extract_spectral_bandwidth(signal_data)
            elif feature == "mfcc":
                feature_dict[feature] = self._extract_mfcc(signal_data)
            else:
                logger.warning(f"未知的频域特征: {feature}")
        
        return feature_dict
    
    def _extract_power_spectral_density(self, signal_data: np.ndarray) -> np.ndarray:
        """
        提取功率谱密度
        
        Args:
            signal_data: 输入信号
            
        Returns:
            功率谱密度
        """
        # FFT变换
        spectrum = fft(signal_data, self.n_fft)
        magnitude_spectrum = np.abs(spectrum[:self.n_fft//2])
        
        # 功率谱密度
        psd = magnitude_spectrum ** 2
        
        return psd
    
    def _extract_spectral_centroid(self, signal_data: np.ndarray) -> float:
        """
        提取谱质心
        
        Args:
            signal_data: 输入信号
            
        Returns:
            谱质心
        """
        # FFT变换
        spectrum = fft(signal_data, self.n_fft)
        magnitude_spectrum = np.abs(spectrum[:self.n_fft//2])
        
        # 频率轴
        freqs = np.linspace(0, self.sample_rate/2, self.n_fft//2)
        
        # 谱质心
        if np.sum(magnitude_spectrum) > 0:
            centroid = np.sum(freqs * magnitude_spectrum) / np.sum(magnitude_spectrum)
        else:
            centroid = 0
        
        return centroid
    
    def _extract_spectral_rolloff(self, signal_data: np.ndarray, roll_percent: float = 0.85) -> float:
        """
        提取谱滚降
        
        Args:
            signal_data: 输入信号
            roll_percent: 滚降百分比
            
        Returns:
            谱滚降频率
        """
        # FFT变换
        spectrum = fft(signal_data, self.n_fft)
        magnitude_spectrum = np.abs(spectrum[:self.n_fft//2])
        
        # 频率轴
        freqs = np.linspace(0, self.sample_rate/2, self.n_fft//2)
        
        # 计算累积能量
        total_energy = np.sum(magnitude_spectrum)
        if total_energy == 0:
            return 0
        
        cumulative_energy = np.cumsum(magnitude_spectrum)
        threshold = roll_percent * total_energy
        
        # 找到滚降点
        rolloff_idx = np.where(cumulative_energy >= threshold)[0]
        if len(rolloff_idx) > 0:
            rolloff_freq = freqs[rolloff_idx[0]]
        else:
            rolloff_freq = freqs[-1]
        
        return rolloff_freq
    
    def _extract_spectral_bandwidth(self, signal_data: np.ndarray) -> float:
        """
        提取谱带宽
        
        Args:
            signal_data: 输入信号
            
        Returns:
            谱带宽
        """
        # 计算谱质心
        centroid = self._extract_spectral_centroid(signal_data)
        
        # FFT变换
        spectrum = fft(signal_data, self.n_fft)
        magnitude_spectrum = np.abs(spectrum[:self.n_fft//2])
        
        # 频率轴
        freqs = np.linspace(0, self.sample_rate/2, self.n_fft//2)
        
        # 谱带宽
        if np.sum(magnitude_spectrum) > 0:
            bandwidth = np.sqrt(np.sum(((freqs - centroid) ** 2) * magnitude_spectrum) / np.sum(magnitude_spectrum))
        else:
            bandwidth = 0
        
        return bandwidth
    
    def _extract_mfcc(self, signal_data: np.ndarray) -> np.ndarray:
        """
        提取MFCC特征
        
        Args:
            signal_data: 输入信号
            
        Returns:
            MFCC特征
        """
        try:
            # 使用librosa提取MFCC
            mfcc = librosa.feature.mfcc(
                y=signal_data.astype(np.float32),
                sr=self.sample_rate,
                n_mfcc=self.n_mfcc,
                n_fft=self.n_fft,
                hop_length=self.hop_length,
                n_mels=self.n_mels
            )
            
            # 取平均值作为特征
            mfcc_mean = np.mean(mfcc, axis=1)
            
            return mfcc_mean
            
        except ImportError:
            logger.warning("librosa未安装，无法提取MFCC特征")
            return np.zeros(self.n_mfcc)
    
    def extract_all_features(self, signal_data: np.ndarray) -> Dict[str, np.ndarray]:
        """
        提取所有频域特征
        
        Args:
            signal_data: 输入信号
            
        Returns:
            所有特征字典
        """
        all_features = [
            "power_spectral_density", "spectral_centroid", 
            "spectral_rolloff", "spectral_bandwidth", "mfcc"
        ]
        
        feature_dict = {}
        for feature in all_features:
            if feature == "power_spectral_density":
                feature_dict[feature] = self._extract_power_spectral_density(signal_data)
            elif feature == "spectral_centroid":
                feature_dict[feature] = self._extract_spectral_centroid(signal_data)
            elif feature == "spectral_rolloff":
                feature_dict[feature] = self._extract_spectral_rolloff(signal_data)
            elif feature == "spectral_bandwidth":
                feature_dict[feature] = self._extract_spectral_bandwidth(signal_data)
            elif feature == "mfcc":
                feature_dict[feature] = self._extract_mfcc(signal_data)
        
        return feature_dict
    
    def get_feature_names(self) -> List[str]:
        """
        获取特征名称列表
        
        Returns:
            特征名称列表
        """
        return self.features.copy()
    
    def get_feature_dimensions(self) -> Dict[str, int]:
        """
        获取特征维度
        
        Returns:
            特征维度字典
        """
        dimensions = {}
        for feature in self.features:
            if feature == "power_spectral_density":
                dimensions[feature] = self.n_fft // 2
            elif feature == "mfcc":
                dimensions[feature] = self.n_mfcc
            else:
                dimensions[feature] = 1
        
        return dimensions
    
    def __str__(self) -> str:
        return f"FrequencyDomainFeatures(features={self.features}, sample_rate={self.sample_rate})"
    
    def __repr__(self) -> str:
        return self.__str__() 